from volcenginesdkarkruntime import Ark
from volcenginesdkarkruntime.types.chat import ChatCompletion, ChatCompletionSystemMessageParam, \
    ChatCompletionUserMessageParam
import json

import function_calling
from basic_config import api_key, knowledge_service_id, system_prompt, model_version
from rag import construct_user_prompt_and_knowledge

# 用 API Key 初始化火山引擎SDK
client = Ark(
    api_key = api_key,
)

messages = []

# 获取用户提出的第一个问题
user_message = ChatCompletionUserMessageParam(role="user", content=input())
if user_message["content"] == "exit":
    exit(0)

user_message = construct_user_prompt_and_knowledge(api_key, knowledge_service_id, user_message)

# 如果存在系统提示词，则添加进会话
if len(system_prompt) > 0:
    system_message = ChatCompletionSystemMessageParam(role="system", content=system_prompt)
    messages.append(system_message)

# 将用户问题正式添加到会话中
messages.append(user_message)

while True:
    # 发起模型请求，由于模型在收到工具执行结果后仍然可能有工具调用意愿，因此需要多次请求
    completion: ChatCompletion = client.chat.completions.create(
        model=model_version,
        messages=messages,
        tools=function_calling.tools
    )
    resp_msg = completion.choices[0].message
    # 展示模型中间过程的回复内容
    print(resp_msg.content)
    if completion.choices[0].finish_reason != "tool_calls":
        # 模型最终总结，没有调用工具意愿
        next_user_message = ChatCompletionUserMessageParam(role="user", content=input())
        if next_user_message.get("content") == "exit":
            break
        messages.append(construct_user_prompt_and_knowledge(api_key, knowledge_service_id, next_user_message))
        continue
    messages.append(completion.choices[0].message.model_dump())
    tool_calls = completion.choices[0].message.tool_calls
    for tool_call in tool_calls:
        tool_name = tool_call.function.name
        if tool_name == "get_current_weather":
            # 调用外部工具
            args = json.loads(tool_call.function.arguments)
            tool_result = function_calling.get_current_weather(**args)
            # 回填工具结果，并获取模型总结回复
            messages.append(
                {"role": "tool", "content": tool_result, "tool_call_id": tool_call.id}
            )